package main

import (
	"context"
	"fmt"
	"log"
	"math/rand"
	"net/url"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

const (
	MaxReconnectAttempts = 10
	HeartbeatInterval    = 30 * time.Second
	HeartbeatTimeout     = 10 * time.Second
)

type WebSocketClient struct {
	url               *url.URL
	conn              *websocket.Conn
	mu                sync.Mutex
	wg                sync.WaitGroup
	isManuallyClosed  bool
	reconnectAttempts int
	ctx               context.Context
	cancel            context.CancelFunc
}

func NewWebSocketClient(urlStr string) (*WebSocketClient, error) {
	u, err := url.Parse(urlStr)
	if err != nil {
		return nil, fmt.Errorf("解析 URL 失败: %w", err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	return &WebSocketClient{
		url:    u,
		ctx:    ctx,
		cancel: cancel,
	}, nil
}

func (c *WebSocketClient) Connect() error {
	c.mu.Lock()
	defer c.mu.Unlock()

	// 检查是否已手动关闭
	if c.isManuallyClosed {
		return fmt.Errorf("客户端已被手动关闭")
	}

	// 清理旧连接
	if c.conn != nil {
		c.conn.Close()
		c.conn = nil
	}

	log.Printf("尝试连接 WebSocket 服务端: %s", c.url.String())
	dialer := websocket.DefaultDialer
	conn, _, err := dialer.Dial(c.url.String(), nil)
	if err != nil {
		return fmt.Errorf("连接失败: %w", err)
	}

	// 🔴 只有连接成功，才进行后续操作
	c.conn = conn
	c.reconnectAttempts = 0 // 重置重连计数

	// 🔴 启动工作 goroutine
	c.wg.Add(2)
	go c.readPump()
	go c.heartbeatPump()

	return nil
}

func (c *WebSocketClient) readPump() {
	defer c.wg.Done()
	defer c.handleConnectionLost() // 连接丢失时触发重连

	for {
		select {
		case <-c.ctx.Done():
			return
		default:
			// 🔴 注意：这里直接使用 c.conn，Connect 成功才启动此 goroutine
			_, message, err := c.conn.ReadMessage()
			if err != nil {
				if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
					log.Printf("读取消息时发生意外错误: %v", err)
				} else {
					log.Printf("连接已关闭 (正常或预期): %v", err)
				}
				return
			}
			log.Printf("收到: %s", message)
		}
	}
}

func (c *WebSocketClient) heartbeatPump() {
	defer c.wg.Done()
	defer c.handleConnectionLost() // 连接丢失时触发重连

	ticker := time.NewTicker(HeartbeatInterval)
	defer ticker.Stop()

	for {
		select {
		case <-c.ctx.Done():
			return
		case <-ticker.C:
			// 🔴 这里直接使用 c.conn
			err := c.conn.WriteMessage(websocket.PingMessage, nil)
			if err != nil {
				log.Printf("发送心跳 (ping) 失败: %v", err)
				return
			}
			c.conn.SetReadDeadline(time.Now().Add(HeartbeatTimeout))
		}
	}
}

func (c *WebSocketClient) handleConnectionLost() {
	c.mu.Lock()
	// 🔴 检查是否已手动关闭或正在处理
	if c.isManuallyClosed || c.conn == nil {
		c.mu.Unlock()
		return
	}

	// 🔴 标记连接已丢失，清理资源
	oldConn := c.conn
	c.conn = nil
	c.mu.Unlock() // 🔴 立即释放锁！非常重要

	// 🔴 关闭旧连接
	if oldConn != nil {
		oldConn.Close()
	}

	// 🔴 启动异步重连
	go c.reconnect()
}

func (c *WebSocketClient) reconnect() {
	for {
		c.mu.Lock()
		// 🔴 检查退出条件
		if c.isManuallyClosed || c.reconnectAttempts >= MaxReconnectAttempts {
			if c.reconnectAttempts >= MaxReconnectAttempts {
				log.Printf("已达到最大重连次数 (%d)，停止重连", MaxReconnectAttempts)
			}
			c.mu.Unlock()
			return
		}
		c.reconnectAttempts++
		delay := calculateBackoffDelay(c.reconnectAttempts)
		jitter := 0.5 + rand.Float64()
		actualDelay := time.Duration(float64(delay) * jitter)
		log.Printf("第 %d 次重连，%v 秒后尝试...", c.reconnectAttempts, actualDelay.Seconds())
		c.mu.Unlock() // 🔴 在等待前释放锁

		// 🔴 等待
		timer := time.NewTimer(actualDelay)
		select {
		case <-timer.C:
			// 🔴 等待结束后，尝试连接
			err := c.Connect()
			if err != nil {
				log.Printf("重连尝试 %d 失败: %v", c.reconnectAttempts, err)
				// 🔴 继续循环进行下一次尝试
				continue
			} else {
				log.Printf("重连成功！")
				return // 🔴 连接成功，退出重连循环
			}
		case <-c.ctx.Done():
			if !timer.Stop() {
				<-timer.C
			}
			return
		}
	}
}

func (c *WebSocketClient) Close() {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.isManuallyClosed {
		return
	}
	c.isManuallyClosed = true
	c.cancel() // 取消 context

	if c.conn != nil {
		c.conn.Close()
		c.conn = nil
	}

	c.wg.Wait() // 等待所有 goroutine 结束
}

func calculateBackoffDelay(attempt int) time.Duration {
	base := 1 * time.Second
	// 指数退避 + 随机抖动
	return base * time.Duration(1<<(attempt-1))
}

func main() {
	rand.Seed(time.Now().UnixNano())

	client, err := NewWebSocketClient("ws://localhost:8080/ws")
	if err != nil {
		log.Fatal("创建客户端失败: ", err)
	}
	defer client.Close()

	// 第一次连接
	if err := client.Connect(); err != nil {
		log.Printf("首次连接失败: %v", err)
		// 🔴 即使首次连接失败，也启动重连机制
		client.reconnect()
	}

	// 保持主 goroutine 运行
	select {}
}
